# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: donggu

from absl import app
from absl import flags
import copy
from experimental.donggu.search.search import Search
import re
import tempfile
import json
import toml

FLAGS = flags.FLAGS
flags.DEFINE_string('target', None, 'BTC-USD.PERPETUAL')
flags.DEFINE_string('dates', None, '20200202-20200205,20200501')
flags.DEFINE_boolean('is_inverse', None, '--is_inverse or --is_inverse=true/false')
flags.DEFINE_boolean('lm_config', None, 'path to lm_config json or toml')
flags.DEFINE_boolean('duration_each_day', '24H', '24H or 3H')

flags.mark_flag_as_required('target')
flags.mark_flag_as_required('dates')
flags.mark_flag_as_required('is_inverse')
flags.mark_flag_as_required('lm_config')

original_config = None
exp_name = None


def set_exp_property(e):
  e.set_property('target', FLAGS.target)
  e.set_property('dates', FLAGS.dates)
  e.set_property('is_inverse', FLAGS.is_inverse)
  e.set_property('lm_config', FLAGS.lm_config)


def set_params_and_cmds(trial):
  global original_config, exp_name
  with tempfile.NamedTemporaryFile('w+', suffix='.json', delete=False) as tmp_config:
    # params
    obj = copy.deepcopy(original_config)
    product_config = obj['products'][FLAGS.target]
    obj['products'] = {FLAGS.target: product_config}

    if FLAGS.is_inverse:
      lotsize = product_config['lot_size']
    else:
      lotsize = 100
    stack = trial.suggest_int('stack', 1, 20)
    threshold = trial.suggest_uniform('threshold', 2, 10)
    lps = trial.suggest_uniform('lean_per_stack_bps', 0, 5)
    horizon = trial.suggest_categorical('horizon', [2, 3, 5, 10, 30])
    pie_cap = trial.suggest_uniform('pie_in_the_sky_cap_bps', 0, 6)

    product_config['lot_size'] = lotsize
    product_config['min_pos'] = -lotsize * stack * 0.5
    product_config['max_pos'] = lotsize * stack * 0.5
    product_config['maintain_open_pos'] = lotsize
    product_config['threshold_bps'] = threshold
    product_config['lean_per_stack_bps'] = lps
    product_config['pie_in_the_sky_cap_bps'] = pie_cap
    product_config['linear_alpha_file'] = f"linear_alpha_zerothres_{horizon}s.json"

    # write config
    json.dump(obj, tmp_config)
    tmp_config.flush()

    cmds = [
        './cc/appcoin2/strategy/lm/run_lm_sim_stdout.sh', exp_name,
        str(trial.number), FLAGS.target, tmp_config.name, FLAGS.dates, FLAGS.duration_each_day
    ]
    return cmds


def main(argv):
  del argv  # Unused

  global original_config, exp_name
  original_config = open(FLAGS.lm_config).read()
  if FLAGS.lm_config.endswith('.toml'):
    original_config = toml.loads(original_config)
  elif FLAGS.lm_config.endswith('.json'):
    original_config = json.loads(original_config)
  else:
    raise RuntimeError()

  exp_name = f'{FLAGS.target}-{FLAGS.dates}'
  Search(
      'lm-agg',
      exp_name,
      set_exp_property,
      set_params_and_cmds,
  ).run()


if __name__ == '__main__':
  app.run(main)
